from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import numpy as np
import pandas as pd
from pykrige.ok import OrdinaryKriging
import json
import io
import base64
from datetime import datetime

app = Flask(__name__)
CORS(app)

# 默认东海区域示例数据
DEFAULT_DATA = {
    "longitude": [122.0, 122.5, 123.0, 123.5, 124.0, 124.5, 125.0, 122.2, 122.8, 123.3, 123.8, 124.3],
    "latitude": [30.0, 30.2, 30.1, 30.3, 30.0, 30.2, 30.1, 29.8, 29.9, 30.0, 30.1, 30.2],
    "temperature": [18.5, 19.2, 18.8, 19.5, 18.9, 19.3, 19.0, 18.7, 19.1, 19.4, 19.6, 19.2]
}

def perform_kriging(lon, lat, temp, grid_size=50):
    """执行克里金插值"""
    # 定义网格范围
    lon_min, lon_max = min(lon) - 0.2, max(lon) + 0.2
    lat_min, lat_max = min(lat) - 0.2, max(lat) + 0.2
    
    # 创建网格
    grid_lon = np.linspace(lon_min, lon_max, grid_size)
    grid_lat = np.linspace(lat_min, lat_max, grid_size)
    
    # 执行普通克里金插值
    OK = OrdinaryKriging(
        lon, lat, temp,
        variogram_model='spherical',
        verbose=False,
        enable_plotting=False,
        nlags=20
    )
    
    # 执行插值
    z, ss = OK.execute('grid', grid_lon, grid_lat)
    
    return grid_lon, grid_lat, z

def create_geojson(grid_lon, grid_lat, z, original_points=None):
    """创建GeoJSON格式的插值结果"""
    features = []
    
    # 添加网格数据（创建多边形网格而不是点）
    for i in range(len(grid_lat) - 1):
        for j in range(len(grid_lon) - 1):
            # 创建网格单元
            coords = [
                [float(grid_lon[j]), float(grid_lat[i])],
                [float(grid_lon[j+1]), float(grid_lat[i])],
                [float(grid_lon[j+1]), float(grid_lat[i+1])],
                [float(grid_lon[j]), float(grid_lat[i+1])],
                [float(grid_lon[j]), float(grid_lat[i])]
            ]
            
            # 计算网格中心温度
            temp = float(z[i, j])
            
            features.append({
                "type": "Feature",
                "geometry": {
                    "type": "Polygon",
                    "coordinates": [coords]
                },
                "properties": {
                    "temperature": temp
                }
            })
    
    # 添加原始采样点
    if original_points:
        for idx, (lon, lat, temp) in enumerate(zip(
            original_points['longitude'], 
            original_points['latitude'], 
            original_points['temperature']
        )):
            features.append({
                "type": "Feature",
                "geometry": {
                    "type": "Point",
                    "coordinates": [float(lon), float(lat)]
                },
                "properties": {
                    "temperature": float(temp),
                    "type": "sample_point",
                    "id": idx
                }
            })
    
    return {
        "type": "FeatureCollection",
        "features": features
    }

@app.route('/')
def index():
    return send_file('../frontend/index.html')

@app.route('/api/interpolate', methods=['POST'])
def interpolate():
    try:
        if 'file' in request.files:
            # 处理上传的文件
            file = request.files['file']
            df = pd.read_csv(file)
            lon = df['longitude'].values
            lat = df['latitude'].values
            temp = df['temperature'].values
        else:
            # 使用默认数据
            lon = np.array(DEFAULT_DATA['longitude'])
            lat = np.array(DEFAULT_DATA['latitude'])
            temp = np.array(DEFAULT_DATA['temperature'])
        
        # 执行克里金插值
        grid_lon, grid_lat, z = perform_kriging(lon, lat, temp)
        
        # 创建GeoJSON
        geojson = create_geojson(
            grid_lon, grid_lat, z,
            {'longitude': lon, 'latitude': lat, 'temperature': temp}
        )
        
        # 计算温度范围
        temp_range = {
            "min": float(np.min(z)),
            "max": float(np.max(z)),
            "mean": float(np.mean(z))
        }
        
        return jsonify({
            "geojson": geojson,
            "temperature_range": temp_range,
            "sample_points": len(lon)
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/export', methods=['POST'])
def export_data():
    try:
        data = request.json
        format_type = data.get('format', 'geojson')
        
        if format_type == 'geojson':
            # 返回GeoJSON文件
            geojson_str = json.dumps(data['geojson'], ensure_ascii=False)
            buffer = io.StringIO(geojson_str)
            return send_file(
                io.BytesIO(buffer.getvalue().encode()),
                mimetype='application/json',
                as_attachment=True,
                download_name=f'sst_interpolation_{datetime.now().strftime("%Y%m%d_%H%M%S")}.geojson'
            )
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
